# -*- coding: utf-8 -*-

from mtl.op_alloc import *
from mtl.op_exec import *
from mtl.op_tran import *
from mtl.op_utils import *

from programs.opcf.conv3d import *
from programs.opcf.fft2d import *


def maru_main():
    # Prepare FFT2D weights.
    fft1d_64_w = fft1d_64_weights()
    addr_fft_w = malloc(len(fft1d_64_w))
    mplace(addr_fft_w, fft1d_64_w)

    # Allocate memory for images.
    channel_size = 4 * 4 * 1024
    image_size = channel_size * 3
    addr_image = malloc(image_size)
    # Set the image value.
    image_data = [0, -16, 0, 0, -4, 30, 2, 4, 0, 8, 0, 0, 0, 16, 0, 0] * 1024
    mplace(addr_image, image_data)
    mplace(addr_image + channel_size, image_data)
    mplace(addr_image + channel_size * 2, image_data)
    # Allocate memory for kernels.
    kernel_size = 4 * 4 * 3
    kernel_data = [2, 3, 1, 2, 3, 4.5, 1.5, 3,
                   1, 1.5, 0.5, 1, 2, 3, 1, 2] * 3
    addr_kernel = malloc(kernel_size)
    mplace(addr_kernel, kernel_data)

    # Conv3D first.
    addr_conv = malloc(64 * 64)
    conv3d_64(
        addrs_image=[addr_image, addr_image + channel_size, addr_image + channel_size * 2],
        addrs_kernel=[addr_kernel, addr_kernel, addr_kernel],
        addr_result=addr_conv)
    # Execute FFT2D
    # fft2d_64(addr_image=addr_conv, addr_weight=addr_fft_w)
    # Then FFT2D for the alogrithm.
    halt()

    print('Memory usage: {}'.format(musage()))
    print('Result saved at: {}'.format(addr_conv))